from typing import Sequence

import numpy as np

import quimb as qu
from quimb.tensor.tensor_1d import TensorNetwork1DOperator, TensorNetwork1DFlat
from quimb.tensor import (
	Tensor, TensorNetwork
)
from quimb.tensor.tensor_core import (
	rand_uuid, make_immutable, oset, tags_to_oset
)
import matplotlib.pyplot as plt

I_o2 = qu.eye(4)
F_o3 = qu.qu([  # [F0, F1, F2, F3]
	[[1.0, 0, 0, 0],
	 [0, 0, 0, 0],
	 [0, 0, 0, 0],
	 [0, 0, 0, 1]],
	[[0, 0, 0, 1],
	 [0, 0, 0, 0],
	 [0, 0, 0, 0],
	 [1, 0, 0, 0]],
	[[0, 0, 0, 0],
	 [0, 1, 0, 0],
	 [0, 0, 1, 0],
	 [0, 0, 0, 0]],
	[[0, 0, 0, 0],
	 [0, 0, 1, 0],
	 [0, -1, 0, 0],
	 [0, 0, 0, 0]]
])
G_o3 = qu.qu([  # [G0, G1, G2, G3]
	[[1.0, 0, 0, 0],
	 [0, 1, 0, 0],
	 [0, 0, 0, 0],
	 [0, 0, 0, 0]],
	[[0, 0, 0, 0],
	 [0, 0, 0, 0],
	 [0, 0, 1, 0],
	 [0, 0, 0, 1]],
	[[0, 1, 0, 0],
	 [1, 0, 0, 0],
	 [0, 0, 0, 0],
	 [0, 0, 0, 0]],
	[[0, 0, 0, 0],
	 [0, 0, 0, 0],
	 [0, 0, 0, 1],
	 [0, 0, -1, 0]]
])
make_immutable(I_o2)
make_immutable(F_o3)
make_immutable(G_o3)

def even_CNOT_mpo(
		L: int,
		tags=None,
		upper_ind_id="k{}",
		lower_ind_id="b{}",
		site_tag_id="I{}",
		**tn_opts,
) -> TensorNetwork:
	sites = tuple(range(L))
	upper_inds, lower_inds, site_tags = \
		map(upper_ind_id.format, sites), map(lower_ind_id.format, sites), map(site_tag_id.format, sites)
	tags = tags_to_oset(tags)

	tn = TensorNetwork1DOperator.new(TensorNetwork1DOperator | TensorNetwork1DFlat,
	                                 L=L,
	                                 upper_ind_id=upper_ind_id,
	                                 lower_ind_id=lower_ind_id,
	                                 site_tag_id=site_tag_id,
	                                 **tn_opts)
	for i in range(L):
		upper_ind, lower_ind, site_tag = next(upper_inds), next(lower_inds), next(site_tags)
		if i % 2:
			t = Tensor(G_o3, inds=(m_ix, upper_ind, lower_ind), tags=tags | oset([site_tag]))
		else:
			if i == L - 1:
				t = Tensor(I_o2, inds=(upper_ind, lower_ind), tags=tags | oset([site_tag]))
			else:
				t = Tensor(F_o3, inds=(m_ix := rand_uuid(), upper_ind, lower_ind), tags=tags | oset([site_tag]))
		tn |= t
	return tn

def odd_CNOT_mpo(
		L: int,
		upper_ind_id="k{}",
		lower_ind_id="b{}",
		site_tag_id="I{}",
		tags=None,
		**tn_opts,
) -> TensorNetwork:
	sites = tuple(range(L))
	upper_inds, lower_inds, site_tags = \
		map(upper_ind_id.format, sites), map(lower_ind_id.format, sites), map(site_tag_id.format, sites)
	tags = tags_to_oset(tags)

	def gen_tensors():
		upper_ind, lower_ind, site_tag = next(upper_inds), next(lower_inds), next(site_tags)
		yield Tensor(I_o2,
		             inds=(upper_ind, lower_ind), tags=tags | oset([site_tag]))
		for i in range(1, L):
			upper_ind, lower_ind, site_tag = next(upper_inds), next(lower_inds), next(site_tags)
			if i % 2:
				if i == L - 1:
					yield Tensor(I_o2, inds=(upper_ind, lower_ind), tags=tags | oset([site_tag]))
				else:
					yield Tensor(F_o3, inds=((m_ix := rand_uuid()), upper_ind, lower_ind), tags=tags | oset([site_tag]))
			else:
				yield Tensor(G_o3, inds=(m_ix, upper_ind, lower_ind), tags=tags | oset([site_tag]))

	return TensorNetwork(gen_tensors(), virtual=True, **tn_opts)

def SPLM_coeff_plot(N_qubit, lam1_mat, lam2_mat,
                    ax=None, title=None, figsize=(14, 3), width=0.2):
	if ax is None:
		_, ax = plt.subplots(figsize=figsize, layout='constrained')

	offset1, offset2 = width * 4 * N_qubit, width * (14 * N_qubit - 10)

	sites1 = np.linspace(0, offset1, 4 * N_qubit)
	sites2 = np.linspace(offset1, offset2, 10 * (N_qubit - 1))

	xticks_labels = ['X', 'Y', 'Z'] * N_qubit + \
	                ['X\nX', 'X\nY', 'X\nZ', 'Y\nX', 'Y\nY', 'Y\nZ', 'Z\nX', 'Z\nY', 'Z\nZ'] * (N_qubit - 1)
	colors1 = ["#7FFFD4", "#76EEC6", "#66CDAA"]
	colors2 = ["#BBFFFF", "#AEEEEE", "#96CDCD", "#668B8B", "#98F5FF", "#8EE5EE", "#7AC5CD", "#53868B", "#00F5FF"]
	for i in range(3):
		ax.bar(sites1[i::4], lam1_mat.T[i], width, color=colors1[i])

	for i in range(9):
		ax.bar(sites2[i::10], lam2_mat.T[i], width, color=colors2[i])

	xticks_sites = np.append(np.delete(sites1, np.arange(3, 4 * N_qubit, 4)),
	                         np.delete(sites2, np.arange(9, 10 * (N_qubit - 1), 10)))

	if title is not None: ax.set_title(title, fontsize=18)
	ax.set_xticks(xticks_sites, xticks_labels)
	ax.set_xlabel("Pauli Operators", fontsize=14)
	ax.set_ylabel(r'Coefficient $(\times 10^{-3})$', fontsize=14)
	ax.set_xlim(- 0.2, width * (14 * N_qubit - 10))
	ax.ticklabel_format(style='sci', scilimits=(0, 0), axis='y')

def pauli_coeff_plot(
		pauli_operators: Sequence[Sequence[str]],
		real_noisy_mpo,
		guess_noisy_mpo,
		ax=None,
		title=None,
		figsize=(8, 5),
		upper_ind_id="k{}",
		lower_ind_id="b{}",
		site_tag_id="I{}",
):
	from states import basic_ptm_vec
	I_o1 = basic_ptm_vec(0, 4).reshape(-1)
	X_o1 = basic_ptm_vec(1, 4).reshape(-1)
	Y_o1 = basic_ptm_vec(2, 4).reshape(-1)
	Z_o1 = basic_ptm_vec(3, 4).reshape(-1)
	pauli_ptm_map = {
		'I': I_o1, 'X': X_o1, 'Y': Y_o1, 'Z': Z_o1,
	}

	if ax is None:
		_, ax = plt.subplots(figsize=figsize, layout='constrained')

	xtick_labels = []
	coeff1s, coeff2s = [], []
	for pauli_operator in pauli_operators:
		mps_u = TensorNetwork([Tensor(pauli_ptm_map[s],
		                              inds=(upper_ind_id.format(l),), tags=[site_tag_id.format(l)])
		                       for l, s in enumerate(pauli_operator)], virtual=True)
		mps_d = TensorNetwork([Tensor(pauli_ptm_map[s],
		                              inds=(lower_ind_id.format(l),), tags=[site_tag_id.format(l)])
		                       for l, s in enumerate(pauli_operator)], virtual=True)
		coeff1 = np.real((mps_u | real_noisy_mpo | mps_d).contract())
		coeff2 = np.real((mps_u | guess_noisy_mpo | mps_d).contract())
		xtick_label = '\n'.join(pauli_operator)
		coeff1s.append(coeff1)
		coeff2s.append(coeff2)
		xtick_labels.append(xtick_label)

	sites = np.arange(0, len(pauli_operators))
	ax.scatter(sites, coeff1s, c='green', marker='o', s=50, edgecolors='k', lw=2,
	           label=r"True noise $\Lambda$")
	ax.scatter(sites, coeff2s, c='orange', marker='x', s=50, edgecolors='k', lw=2,
	           label=r"Restructured $\Lambda_\theta$")
	if title is not None: ax.set_title(title, fontsize=18)
	ax.set_xticks(sites, xtick_labels)
	ax.set_xlabel("Pauli Operators", fontsize=14)
	ax.set_ylabel(r'Coefficient', fontsize=14)
	ax.set_xlim([-0.5, len(pauli_operators) + 0.5])
	ax.legend()
